import pymysql

class MySQLDB:
    def __init__(self, db_name, user='root', password='rootroot', host='localhost', port=3306):
        self.db_name = db_name
        self.conn = pymysql.connect(
            host=host,
            user=user,
            password=password,
            port=port,
            charset='utf8mb4',
            cursorclass=pymysql.cursors.DictCursor
        )
        self.cursor = self.conn.cursor()
        
        # 检查并创建数据库
        #self.create_database_if_not_exists(db_name)
        
        # 连接到指定的数据库
        self.conn.select_db(db_name)

    def create_database_if_not_exists(self, db_name):
        self.cursor.execute(f"CREATE DATABASE IF NOT EXISTS `{db_name}` DEFAULT CHARACTER SET utf8mb4")
        self.conn.commit()

    def table_exists(self, table_name):
        self.cursor.execute("SHOW TABLES LIKE %s", (table_name,))
        result = self.cursor.fetchone()
        
        return result is not None



    def create_table(self, tb_name, zdnames: dict, uniquezd=[], default=[]):
        zdlist = []
        print(f'数据库{self.db_name}不存在表{tb_name},新建表《{tb_name}》')
        for k, v in zdnames.items():
            lx = 'VARCHAR(255)' if isinstance(v, str) else 'INT'
            if v in ['DATE','DATETIME','TIMESTAMP']:
                lx=v

            if k in default:
                zdlist.append(f'{k} INT DEFAULT 0')
            else:
                zdlist.append(f'{k} {lx} NULL')
            

        for uni_item in uniquezd:
            key_str=','.join(uni_item)
            zdlist.append(f'UNIQUE KEY ({key_str})')

        zdstr = ','.join(zdlist)
        self.cursor.execute(f'''
            CREATE TABLE IF NOT EXISTS {tb_name} (
                id INT PRIMARY KEY AUTO_INCREMENT,
                {zdstr},
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
            )
        ''')
        self.conn.commit()

    def add_columns_if_not_exists(self, table_name, columns: dict):
        for column_name, column_definition in columns.items():
            self.cursor.execute("""
                SELECT COUNT(*) 
                FROM INFORMATION_SCHEMA.COLUMNS 
                WHERE TABLE_NAME = %s 
                AND COLUMN_NAME = %s
            """, (table_name, column_name))
            if self.cursor.fetchone()['COUNT(*)'] == 0:
                self.cursor.execute(f"ALTER TABLE {table_name} ADD COLUMN {column_name} {column_definition}")
                self.conn.commit()

    def get_all_tables(self):

        try:

            # 执行查询表名的 SQL
            self.cursor.execute("SHOW TABLES")
            tables = self.cursor.fetchall()  # 获取所有表名
            return [list(table.values())[0] for table in tables]  # 返回表名列表
        except pymysql.MySQLError as e:
            print(f"Error while fetching tables: {e}")
            return []


    def get_table_columns(self, table_name):
        """
        获取指定表的列名
        :param table_name: 表名
        :return: 包含列名的列表
        """
        try:
            query = f"SHOW COLUMNS FROM `{table_name}`"
            self.cursor.execute(query)
            columns = [row['Field'] for row in self.cursor.fetchall()]  # 第一列为列名
            return columns
        except pymysql.MySQLError as e:
            print(f"Error fetching columns: {e}")
            return []


    def query(self, sqlstr, params=(), return_as_dict=False):
        self.cursor.execute(sqlstr, params)
        rows = self.cursor.fetchall()
        if return_as_dict:
            return rows  # 返回字典数组
        else:
            return [tuple(row.values()) for row in rows]  # 返回元组数组

    def execute_command(self, sqlstr, params=()):
        try:
            self.cursor.execute(sqlstr, params)
            self.conn.commit()
            return self.cursor.rowcount
        except Exception as e:
            print(e)
            return 0

    def close_connection(self):
        self.cursor.close()
        self.conn.close()
